#include <algorithm>
#include <string>

#include "my_assert.h"
#include "memory_storage.h"
#include "../logger/logger.h"

using std::to_string;
using std::string;

MemoryStorage::MemoryStorage(uint64_t size, bool addresses_are_virtual):
	chunk_addrs_(),
	chunk_sizes_(),
	allocated_(),
	memory_(nullptr),
	size_(size),
	addresses_are_virtual_(addresses_are_virtual) {
	LOG_FUNCTION("");	

	if (!addresses_are_virtual) 
		memory_ = new char[size];

	MemoryStorageChunk *pchunk = new MemoryStorageChunk;
	pchunk->size_ = size;
	pchunk->mem_ = memory_;
	chunk_addrs_.insert(pchunk);
	chunk_sizes_.insert(pchunk);
	AssertOk();
}

MemoryStorage::~MemoryStorage() {
	LOG_FUNCTION("");	
	AssertOk();
	my_assert(chunk_addrs_.size() == chunk_sizes_.size());
	for (auto pchunk : allocated_) {
		my_assert(pchunk);
		delete pchunk;
	}
	if (!addresses_are_virtual_) {
		my_assert(memory_);
		delete [] memory_;
	}
}

StorageChunk *MemoryStorage::Alloc(uint64_t size,
    const std::list<StorageChunk**>& currently_unused,
    std::function<void (StorageChunk**)> delete_func) {
	LOG_FUNCTION("");	
	AssertOk();
	LOGGER_WRITE(Logger::Priority::INFO, "size = " + 
	    to_string(size));
	uint64_t arranged_size = size;
	if (size % ARRANGEMENT_SIZE)
		size += ARRANGEMENT_SIZE - size % ARRANGEMENT_SIZE;
	MemoryStorageChunk to_find_chunk; 
	to_find_chunk.size_ = arranged_size - 1; // because upper_bound(n) == n + 1
	LOGGER_WRITE(Logger::Priority::INFO, "arranged_size_ = " + 
	    to_string(arranged_size));

	ChunkSizesSet::iterator ipempty_chunk = 
	    chunk_sizes_.upper_bound(&to_find_chunk);
	for (auto pchunk : chunk_sizes_) {
		LOGGER_WRITE(Logger::Priority::INFO, "pchunk.size() = " + 
		    to_string(pchunk->size_));
	}
	
	while (ipempty_chunk == chunk_sizes_.end()) {
		if (currently_unused.empty()) {
			LOGGER_WRITE(Logger::Priority::INFO, "return nullptr");
			return nullptr;
		}
		delete_func(currently_unused.front());
		ipempty_chunk = chunk_sizes_.upper_bound(&to_find_chunk);
		// delete_func will modify currently_unused list
	}
	
	MemoryStorageChunk *pempty_chunk = *ipempty_chunk;
	my_assert(pempty_chunk);
	my_assert(pempty_chunk->size_ >= arranged_size);
	EraseBySizesIt(ipempty_chunk);
	if (pempty_chunk->size_ == arranged_size) {
		allocated_.insert(pempty_chunk);
		return pempty_chunk;
	}

	MemoryStorageChunk *pnew_chunk = new MemoryStorageChunk;
	pnew_chunk->mem_ = pempty_chunk->mem_;
	pnew_chunk->size_ = arranged_size;

	pempty_chunk->mem_ += pnew_chunk->size_;
	pempty_chunk->size_ -= pnew_chunk->size_;

	chunk_sizes_.insert(pempty_chunk);
	chunk_addrs_.insert(pempty_chunk);

	allocated_.insert(pnew_chunk);
	AssertOk();
	my_assert(addresses_are_virtual_ || pnew_chunk->mem_);
	return pnew_chunk;
}

void *MemoryStorage::Load(StorageChunk *p_st_chunk){
	LOG_FUNCTION("");	
	AssertOk();
	my_assert(p_st_chunk);
	my_assert(allocated_.find(p_st_chunk) != allocated_.end());
	
	AssertOk();
	return ((MemoryStorageChunk*)p_st_chunk)->mem_;	

}

void MemoryStorage::Free(StorageChunk *p_st_chunk) {
	LOG_FUNCTION("");	
	AssertOk();
	my_assert(p_st_chunk);
	auto ipfree_chunk = allocated_.find(p_st_chunk);
	my_assert(ipfree_chunk != allocated_.end());
	MemoryStorageChunk *pfree_chunk = (MemoryStorageChunk*) *ipfree_chunk;
	allocated_.erase(ipfree_chunk);
	my_assert(pfree_chunk);

	auto iplower_chunk = chunk_addrs_.lower_bound(pfree_chunk);
	auto ipupper_chunk = chunk_addrs_.upper_bound(pfree_chunk);
	MemoryStorageChunk *plower_chunk = nullptr, *pupper_chunk = nullptr;
	bool lower_changed = false, upper_changed = false;
	
	if (iplower_chunk != chunk_addrs_.end()) {
		plower_chunk = *iplower_chunk;
		my_assert(plower_chunk);
		if (plower_chunk->mem_ + plower_chunk->size_ == 
		    pfree_chunk->mem_) {
			EraseByAddrsIt(iplower_chunk);
			plower_chunk->size_ += pfree_chunk->size_;
			lower_changed = true;
		}
	}

	if (ipupper_chunk != chunk_addrs_.end()) {
		pupper_chunk = *ipupper_chunk;
		my_assert(pupper_chunk);
		if (pfree_chunk->mem_ + pfree_chunk->size_ == 
		    (*ipupper_chunk)->mem_) {
			if (lower_changed) {
				plower_chunk->size_ += pupper_chunk->size_;
				EraseByAddrsIt(ipupper_chunk);
				delete pupper_chunk;
			}
			else {
				EraseByAddrsIt(ipupper_chunk);
				pupper_chunk->mem_ -= pfree_chunk->size_;
				pupper_chunk->size_ += pfree_chunk->size_;
				upper_changed = true;
			}
		}
	}
	LOGGER_WRITE(Logger::Priority::INFO, "addrs.size() = " + 
	    to_string(chunk_addrs_.size()) + ", sizes.size() = " + 
	    to_string(chunk_sizes_.size()));

	bool free_chunk_needs_delete = true;
	if (lower_changed) {
		LOGGER_WRITE(Logger::Priority::INFO, "lower_changed");
		chunk_addrs_.insert(plower_chunk);
		chunk_sizes_.insert(plower_chunk);
	} else if (upper_changed) {
		LOGGER_WRITE(Logger::Priority::INFO, "upper_changed");
		LOGGER_WRITE(Logger::Priority::INFO, "upper->size = " + 
		    to_string(pupper_chunk->size_));
		chunk_addrs_.insert(pupper_chunk);
		chunk_sizes_.insert(pupper_chunk);
	} else {
		LOGGER_WRITE(Logger::Priority::INFO, "free_added");
		chunk_addrs_.insert(pfree_chunk);
		chunk_sizes_.insert(pfree_chunk);
		free_chunk_needs_delete = false;		
	}

	LOGGER_WRITE(Logger::Priority::INFO, "addrs.size() = " + 
	    to_string(chunk_addrs_.size()) + ", sizes.size() = " + 
	    to_string(chunk_sizes_.size()));

	if (free_chunk_needs_delete)
		delete pfree_chunk;
	AssertOk();
}

void MemoryStorage::AssertOk() {
	LOGGER_WRITE(Logger::Priority::DEBUG, "addrs.size() = " + 
	    to_string(chunk_addrs_.size()) + ", sizes.size() = " + 
	    to_string(chunk_sizes_.size()));
	my_assert(chunk_addrs_.size() == chunk_sizes_.size());
	if (!addresses_are_virtual_)
		my_assert(memory_);
	my_assert(size_ > 0);
	
	uint64_t size_sum1 = 0;
	for (auto pmem_chunk : chunk_sizes_) {
		my_assert(pmem_chunk);
		my_assert(pmem_chunk->size_ > 0);
		if (!addresses_are_virtual_)
			my_assert(pmem_chunk->mem_);
		size_sum1 += pmem_chunk->size_;
	}

	uint64_t size_sum2 = 0;
	for (auto pmem_chunk : chunk_addrs_) {
		my_assert(pmem_chunk);
		my_assert(pmem_chunk->size_ > 0);
		if (!addresses_are_virtual_)
			my_assert(pmem_chunk->mem_);
		size_sum2 += pmem_chunk->size_;
	}
	my_assert(size_sum1 == size_sum2);
	
	uint64_t size_sum = size_sum1;
	for (auto pmem_chunk : allocated_) {
		my_assert(pmem_chunk);
		my_assert(((MemoryStorageChunk*)pmem_chunk)->size_ > 0);
		size_sum += ((MemoryStorageChunk*)pmem_chunk)->size_;
	}

	LOGGER_WRITE(Logger::Priority::INFO, "size_sum = " + 
	    to_string(size_sum) + ", size_ = " + to_string(size_));
	my_assert(size_sum == size_);
}

void MemoryStorage::EraseBySizesIt(ChunkSizesSet::iterator i) {
	LOG_FUNCTION("");	
	my_assert(i != chunk_sizes_.end());
	MemoryStorageChunk *pchunk = *i;
	chunk_sizes_.erase(i);
	my_assert(pchunk);
	auto addrs_erase_it = chunk_addrs_.equal_range(pchunk);
	my_assert(addrs_erase_it.first != chunk_addrs_.end());
	for (auto it = addrs_erase_it.first; it != addrs_erase_it.second; ++it)
		if (*it == pchunk) {
			chunk_addrs_.erase(it);
			break;
		}
}

void MemoryStorage::EraseByAddrsIt(ChunkAddrsSet::iterator i) {
	LOG_FUNCTION("");	
	my_assert(i != chunk_addrs_.end());
	MemoryStorageChunk *pchunk = *i;
	chunk_addrs_.erase(i);
	my_assert(pchunk);
	auto sizes_erase_it = chunk_sizes_.equal_range(pchunk);
	my_assert(sizes_erase_it.first != chunk_sizes_.end());
	for (auto it = sizes_erase_it.first; it != sizes_erase_it.second; ++it)
		if (*it == pchunk) {
			chunk_sizes_.erase(it);
			break;
		}
}
